package com.wt.admin.service.ai.impl.mcp;

import com.wt.admin.domain.entity.ai.MCPEntity;
import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.mcp.AsyncMcpToolCallbackProvider;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.mcp.client.autoconfigure.NamedClientMcpTransport;
import org.springframework.ai.mcp.client.autoconfigure.configurer.McpAsyncClientConfigurer;
import org.springframework.ai.mcp.client.autoconfigure.configurer.McpSyncClientConfigurer;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpClientCommonProperties;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpSseClientProperties;
import org.springframework.ai.mcp.customizer.McpAsyncClientCustomizer;
import org.springframework.ai.mcp.customizer.McpSyncClientCustomizer;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;


/**
 * 拆解spring boot加载MCP的配置文件
 */
@Component
@EnableConfigurationProperties({McpClientCommonProperties.class})
@Slf4j
public class MCPStart {

    @Resource
    McpClientCommonProperties properties;

    @Resource
    MCPTransport mcpTransport;

    /**
     * stdio
     * @return ToolCallbackProvider
     * @throws Exception
     */
    @SneakyThrows
    public ToolCallbackProvider builderMcpClientSTDIO(MCPEntity mcpEntities) {
        List<NamedClientMcpTransport> namedClientMcpTransports = mcpTransport.stdioTransports(mcpEntities.getContent());
        return getToolCallbackProvider(namedClientMcpTransports, mcpEntities.isSync());
    }

    /**
     * sse
     * @return ToolCallbackProvider
     */
    public ToolCallbackProvider builderMcpClientSSE(MCPEntity mcpEntities) {
        List<NamedClientMcpTransport> namedClientMcpTransports = mcpTransport.mcpHttpClientTransports(
                new HashMap<String, McpSseClientProperties.SseParameters>() {{
                    put(mcpEntities.getTitle(), new McpSseClientProperties.SseParameters(mcpEntities.getAddress(),mcpEntities.getSseEndpoint()));
                }}
        );
        return getToolCallbackProvider(namedClientMcpTransports, mcpEntities.isSync());
    }

    /**
     * Flux
     * @return ToolCallbackProvider
     */
    public ToolCallbackProvider builderMcpClientFlux(MCPEntity mcpEntities) {
        List<NamedClientMcpTransport> namedClientMcpTransports = mcpTransport.webFluxClientTransports(
                new HashMap<String, McpSseClientProperties.SseParameters>() {{
                    put(mcpEntities.getTitle(), new McpSseClientProperties.SseParameters(mcpEntities.getAddress(),mcpEntities.getSseEndpoint()));
                }}
        );
        return getToolCallbackProvider(namedClientMcpTransports, mcpEntities.isSync());
    }

    private ToolCallbackProvider getToolCallbackProvider(List<NamedClientMcpTransport> namedTransports, boolean sync) {
        if (sync) {
            List<McpSyncClient> mcpSyncClients = mcpSyncClients(namedTransports);
            return getSyncToolCallbackProvider(mcpSyncClients);
        } else {
            List<McpAsyncClient> mcpSyncClients = mcpAsyncClients(namedTransports);
            return getAsyncToolCallbackProvider(mcpSyncClients);
        }
    }

    /**
     * 同步方式
     *
     * @param namedTransports
     * @return
     */
    private List<McpSyncClient> mcpSyncClients(List<NamedClientMcpTransport> namedTransports) {
        List<McpSyncClientCustomizer> mcpSyncClientCustomizers = new ArrayList<>();
        McpSyncClientConfigurer mcpSyncClientConfigurer = new McpSyncClientConfigurer(mcpSyncClientCustomizers);
        List<McpSyncClient> mcpSyncClients = new ArrayList<>();
        McpSyncClient syncClient;
        if (!CollectionUtils.isEmpty(namedTransports)) {
            for (Iterator<NamedClientMcpTransport> var6 = namedTransports.iterator(); var6.hasNext(); mcpSyncClients.add(syncClient)) {
                NamedClientMcpTransport namedTransport = var6.next();

                McpSchema.Implementation clientInfo = new McpSchema.Implementation(
                        this.connectedClientName(properties.getName(), namedTransport.name()),
                        properties.getVersion());

                McpClient.SyncSpec syncSpec = McpClient.sync(namedTransport.transport())
                        .clientInfo(clientInfo).requestTimeout(properties.getRequestTimeout());

                syncSpec = mcpSyncClientConfigurer.configure(namedTransport.name(), syncSpec);
                syncClient = syncSpec.build();
                if (properties.isInitialized()) {
                    syncClient.initialize();
                }
            }
        }
        return mcpSyncClients;
    }


    /**
     * 异步方式
     *
     * @param namedTransports
     * @return
     */
    private List<McpAsyncClient> mcpAsyncClients(List<NamedClientMcpTransport> namedTransports) {
        List<McpAsyncClientCustomizer> mcpSyncClientCustomizers = new ArrayList<>();
        McpAsyncClientConfigurer mcpSyncClientConfigurer = new McpAsyncClientConfigurer(mcpSyncClientCustomizers);

        List<McpAsyncClient> mcpSyncClients = new ArrayList<>();
        McpAsyncClient syncClient;
        if (!CollectionUtils.isEmpty(namedTransports)) {
            for (Iterator<NamedClientMcpTransport> var6 = namedTransports.iterator(); var6.hasNext(); mcpSyncClients.add(syncClient)) {
                NamedClientMcpTransport namedTransport = var6.next();
                McpSchema.Implementation clientInfo = new McpSchema.Implementation(
                        this.connectedClientName(properties.getName(), namedTransport.name()),
                        properties.getVersion());

                McpClient.AsyncSpec syncSpec = McpClient.async(namedTransport.transport())
                        .clientInfo(clientInfo).requestTimeout(properties.getRequestTimeout());

                syncSpec = mcpSyncClientConfigurer.configure(namedTransport.name(), syncSpec);
                syncClient = syncSpec.build();
                if (properties.isInitialized()) {
                    syncClient.initialize().block();
                }
            }
        }
        return mcpSyncClients;
    }


    /**
     * 同步方式
     *
     * @param mcpSyncClients
     */
    private SyncMcpToolCallbackProvider getSyncToolCallbackProvider(List<McpSyncClient> mcpSyncClients) {
        SyncMcpToolCallbackProvider toolCallbackProvider = new SyncMcpToolCallbackProvider(mcpSyncClients);
        for (McpSyncClient toolCallback : mcpSyncClients) {
            log.debug("同步,成功加载MCP工具 {} -> {} ,有如下方法：", toolCallback.getServerInfo().name(), toolCallback.getServerInfo().version());
            for (ToolCallback callback : toolCallbackProvider.getToolCallbacks()) {
                ToolDefinition toolDefinition = callback.getToolDefinition();
                log.debug("{}：{} ",toolDefinition.name(), toolDefinition.description());
            }
        }
        return toolCallbackProvider;
    }

    /**
     * 异步方式
     *
     * @param mcpSyncClients
     * @return
     */
    private AsyncMcpToolCallbackProvider getAsyncToolCallbackProvider(List<McpAsyncClient> mcpSyncClients) {
        AsyncMcpToolCallbackProvider toolCallbackProvider = new AsyncMcpToolCallbackProvider(mcpSyncClients);
        for (McpAsyncClient toolCallback : mcpSyncClients) {
            log.debug("异步,成功加载MCP工具 {} -> {} ,有如下方法：", toolCallback.getServerInfo().name(), toolCallback.getServerInfo().version());
            for (ToolCallback callback : toolCallbackProvider.getToolCallbacks()) {
                ToolDefinition toolDefinition = callback.getToolDefinition();
                log.debug("{}：{} ",toolDefinition.name(), toolDefinition.description());
            }
        }
        return toolCallbackProvider;
    }

    private String connectedClientName(String clientName, String serverConnectionName) {
        return clientName + "-" + serverConnectionName;
    }

}
